package main

import (
	"SQLi/model"
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"log"
	"strings"
)

type UserRepository struct {
	db *sql.DB
}

func NewUserRepository(db *sql.DB) *UserRepository {
	return &UserRepository{db: db}
}

func GetDB(username, password string) (*sql.DB, error) {
	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(localhost:3306)/go_sec_labs?interpolateParams=true", username, password))
	if err != nil {
		return nil, err
	}

	// Test the connection
	if err := db.Ping(); err != nil {
		return nil, err
	}
	return db, nil
}

func (repo *UserRepository) SelectByName(name string) ([]model.User, error) {
	//query := fmt.Sprintf("SELECT id, name, age, sex, password FROM user WHERE name = '%s'", name)	//漏洞代码
	query := "SELECT id, name, age, sex, password FROM user WHERE name = ?"
	rows, err := repo.db.Query(query, name)
	if err != nil {
		log.Println(err)
		return nil, err
	}
	defer rows.Close()

	return repo.orm2User(rows)
}

func (repo *UserRepository) SelectIn(names []string) ([]model.User, error) {
	// 构建in查询字符串 "?, ?" 形式
	str, args := func(names []string) (string, []interface{}) {
		placeholders := make([]string, len(names))
		args := make([]interface{}, len(names))
		for i, id := range names {
			placeholders[i] = "?"
			args[i] = id
		}
		return strings.Join(placeholders, ","), args
	}(names)

	//漏洞代码
	//str := "'" + strings.Join(names, "','") + "'"
	query := fmt.Sprintf("SELECT id, name, age, sex, password FROM user WHERE name in (%s)", str)
	rows, err := repo.db.Query(query, args...)
	if err != nil {
		log.Println(err)
		return nil, err
	}
	defer rows.Close()

	return repo.orm2User(rows)
}

func (repo *UserRepository) SelectLike(name string) ([]model.User, error) {
	//漏洞代码
	//query := "SELECT id, name, age, sex, password FROM user WHERE name like '%" + name + "%'"

	query := "SELECT id, name, age, sex, password FROM user WHERE name like ?"
	rows, err := repo.db.Query(query, "%"+name+"%")
	if err != nil {
		log.Println(err)
		return nil, err
	}
	defer rows.Close()

	return repo.orm2User(rows)
}

func (repo *UserRepository) SelectOrderBy(orderField string) ([]model.User, error) {
	// 定义允许的排序字段白名单
	allowedFields := map[string]bool{
		"id":    true,
		"name":  true,
		"age":   true,
		"sex":   true,
	}
	// 检查用户提供的排序字段是否在白名单中
	if !allowedFields[orderField] {
		log.Println("Invalid order field:", orderField)
		return nil, fmt.Errorf("invalid order field: %s", orderField)
	}

	//漏洞代码
	query := fmt.Sprintf("SELECT id, name, age, sex, password FROM user order by %s desc", orderField)

	rows, err := repo.db.Query(query)
	if err != nil {
		log.Println(err)
		return nil, err
	}
	defer rows.Close()

	return repo.orm2User(rows)
}

func (repo *UserRepository) SelectPage(limit int, offset int) ([]model.User, error) {
	//漏洞代码

	query := fmt.Sprintf("SELECT id, name, age, sex, password FROM user limit ? offset ?")
	rows, err := repo.db.Query(query, limit, offset)
	if err != nil {
		log.Println(err)
		return nil, err
	}
	defer rows.Close()

	return repo.orm2User(rows)
}

func (repo *UserRepository) Insert(users []model.User) (int64, error) {
	// 漏洞代码
	//query := "INSERT INTO `user` (`name`, `age`, `sex`, `password`) VALUES "
	//for _, user := range users {
	//	query += fmt.Sprintf("('%s',%d,'%s','%s'),", user.Name, user.Age, user.Sex, user.Password)
	//}
	//query = query[0:len(query)-1]

	// 构建占位符部分
	placeholders := make([]string, len(users))
	args := make([]interface{}, 0, len(users)*4)
	for i, user := range users {
		placeholders[i] = "(?, ?, ?, ?)"
		args = append(args, user.Name, user.Age, user.Sex, user.Password)
	}

	query := fmt.Sprintf("INSERT INTO `user` (`name`, `age`, `sex`, `password`) VALUES %s", strings.Join(placeholders, ","))

	result, err := repo.db.Exec(query, args...)
	if err != nil {
		log.Println(err)
		return 0, err
	}

	rowsAffected, err := result.RowsAffected()
	if err!= nil {
		log.Println(err)
		return 0, err
	}

	return rowsAffected, nil
}

func (repo *UserRepository) orm2User(rows *sql.Rows) ([]model.User, error) {
	var users []model.User
	for rows.Next() {
		var user model.User
		if err := rows.Scan(&user.Id, &user.Name, &user.Age, &user.Sex, &user.Password); err != nil {
			log.Printf("Failed to scan row: %v", err)
			continue
		}
		users = append(users, user)
	}

	if err := rows.Err(); err != nil {
		log.Printf("Row iteration error: %v", err)
		return nil, err
	}

	return users, nil
}

func main() {
	// 初始化数据库连接
	db, err := GetDB("root", "111111")
	if err != nil {
		log.Fatalf("Get db error: %v", err)
	}
	defer db.Close() // 确保程序结束时关闭数据库连接

	userRepo := NewUserRepository(db)

	//users, err := userRepo.SelectByName("' or 1=1 -- ")
	//users, err := userRepo.SelectIn([]string{"zhangsan", "wangwu') or 1=1 -- "})
	//users, err := userRepo.SelectLike("an%' or 1=1 -- ")
	//users, err := userRepo.SelectOrderBy("if(substr(database(),1,1)='a',id,age)")
	//users, err := userRepo.SelectPage(2, 3)

	id, _ := userRepo.Insert([]model.User{
		{Name: "'+if(substr(database(),1,1)='g',sleep(3),1)+'", Age: 44, Sex: "0", Password: "12121"},
		{Name: "bb", Age: 66, Sex: "1", Password: "32323"},
	})
	fmt.Printf("rowsAffected: %v\n", id)

	//if err != nil {
	//	log.Fatalf("SelectByName error: %v", err)
	//}
	//for _, user := range users {
	//	fmt.Println(user)
	//}
}